#!/usr/bin/env python3
"""
提示融合模块

融合记忆模块、周期模块、质量控制模块的输出，生成最优掩码提示
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Tuple, Optional, Any

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from utils.logger import get_logger


class PromptFusion(nn.Module):
    """提示融合模块"""
    
    def __init__(self,
                 prompt_dim: int = 256,
                 fusion_method: str = "attention",  # "attention", "weighted", "concat"
                 device: str = "cuda"):
        """
        初始化提示融合模块
        
        Args:
            prompt_dim: 提示特征维度
            fusion_method: 融合方法
            device: 设备
        """
        super().__init__()
        
        self.prompt_dim = prompt_dim
        self.fusion_method = fusion_method
        self.device = device
        self.logger = get_logger("PromptFusion")
        
        if fusion_method == "attention":
            # 注意力融合
            self.fusion_attention = nn.MultiheadAttention(
                embed_dim=prompt_dim,
                num_heads=8,
                dropout=0.1,
                batch_first=True
            )
            self.fusion_projection = nn.Sequential(
                nn.Linear(prompt_dim, prompt_dim),
                nn.ReLU(),
                nn.Linear(prompt_dim, prompt_dim)
            )
        elif fusion_method == "weighted":
            # 加权融合
            self.fusion_weights = nn.Parameter(torch.ones(3) / 3)  # 三个模块
            self.fusion_projection = nn.Sequential(
                nn.Linear(prompt_dim, prompt_dim),
                nn.ReLU(),
                nn.Linear(prompt_dim, prompt_dim)
            )
        elif fusion_method == "concat":
            # 拼接融合
            self.fusion_projection = nn.Sequential(
                nn.Linear(prompt_dim * 3, prompt_dim * 2),
                nn.ReLU(),
                nn.Linear(prompt_dim * 2, prompt_dim)
            )
        
        # 掩码提示生成器（从融合特征生成掩码提示）
        self.mask_prompt_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 32 * 32)  # 32x32掩码
        )
        
        # 置信度评估器
        self.confidence_estimator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 1),
            nn.Sigmoid()
        )
    
    def fuse_features(self,
                     memory_features: Optional[torch.Tensor],
                     periodic_features: Optional[torch.Tensor],
                     quality_features: Optional[torch.Tensor]) -> torch.Tensor:
        """
        融合三个模块的特征
        
        Args:
            memory_features: 记忆模块特征 (B, prompt_dim) 或 (B, H, W, prompt_dim)
            periodic_features: 周期模块特征 (B, prompt_dim) 或 (B, H, W, prompt_dim)
            quality_features: 质量模块特征 (B, prompt_dim) 或 (B, H, W, prompt_dim)
            
        Returns:
            融合后的特征 (B, prompt_dim) 或 (B, H, W, prompt_dim)
        """
        # 收集所有非None的特征
        features_list = []
        if memory_features is not None:
            features_list.append(memory_features)
        if periodic_features is not None:
            features_list.append(periodic_features)
        if quality_features is not None:
            features_list.append(quality_features)
        
        if len(features_list) == 0:
            # 如果没有特征，返回零特征
            batch_size = 1
            if memory_features is not None:
                batch_size = memory_features.shape[0]
            elif periodic_features is not None:
                batch_size = periodic_features.shape[0]
            elif quality_features is not None:
                batch_size = quality_features.shape[0]
            
            return torch.zeros(batch_size, self.prompt_dim, device=self.device)
        
        # 统一特征形状
        # 如果是空间特征，先展平
        if len(features_list[0].shape) > 2:
            # 空间特征 (B, H, W, C) -> (B, H*W, C)
            features_list = [f.view(f.shape[0], -1, f.shape[-1]) for f in features_list]
        
        # 统一到相同维度
        features_list = [self._normalize_feature(f) for f in features_list]
        
        if self.fusion_method == "attention":
            # 注意力融合
            # 堆叠特征 (B, num_modules, prompt_dim)
            stacked_features = torch.stack(features_list, dim=1)
            
            # 自注意力
            fused, _ = self.fusion_attention(
                stacked_features, stacked_features, stacked_features
            )
            
            # 平均池化
            fused = torch.mean(fused, dim=1)  # (B, prompt_dim)
            
            # 投影
            fused = self.fusion_projection(fused)
            
        elif self.fusion_method == "weighted":
            # 加权融合
            weights = F.softmax(self.fusion_weights[:len(features_list)], dim=0)
            
            fused = sum(w * f for w, f in zip(weights, features_list))
            fused = torch.mean(fused, dim=1) if len(fused.shape) > 2 else fused
            fused = self.fusion_projection(fused)
            
        elif self.fusion_method == "concat":
            # 拼接融合
            if len(features_list) == 1:
                concat_features = features_list[0]
            else:
                concat_features = torch.cat(features_list, dim=-1)
            
            fused = self.fusion_projection(concat_features)
            fused = torch.mean(fused, dim=1) if len(fused.shape) > 2 else fused
        
        return fused
    
    def _normalize_feature(self, feature: torch.Tensor) -> torch.Tensor:
        """归一化特征到prompt_dim维度"""
        if feature.shape[-1] != self.prompt_dim:
            # 如果维度不匹配，使用线性投影
            if not hasattr(self, f'_proj_{id(feature)}'):
                setattr(self, f'_proj_{id(feature)}', 
                       nn.Linear(feature.shape[-1], self.prompt_dim).to(self.device))
            proj = getattr(self, f'_proj_{id(feature)}')
            feature = proj(feature)
        return feature
    
    def generate_mask_prompt(self,
                            fused_features: torch.Tensor,
                            image_shape: Tuple[int, int]) -> torch.Tensor:
        """
        从融合特征生成掩码提示
        
        Args:
            fused_features: 融合后的特征 (B, prompt_dim)
            image_shape: 图像形状 (H, W)
            
        Returns:
            掩码提示 (B, 1, 32, 32)
        """
        batch_size = fused_features.shape[0]
        
        # 生成掩码logits
        mask_logits = self.mask_prompt_generator(fused_features)  # (B, 32*32)
        
        # 重塑为2D掩码
        mask_prompt = mask_logits.view(batch_size, 32, 32)
        mask_prompt = mask_prompt.unsqueeze(1)  # (B, 1, 32, 32)
        
        # 如果需要，可以上采样到原始图像尺寸
        if image_shape != (32, 32):
            mask_prompt = F.interpolate(
                mask_prompt,
                size=image_shape,
                mode='bilinear',
                align_corners=False
            )
        
        return mask_prompt
    
    def compute_confidence(self, fused_features: torch.Tensor) -> torch.Tensor:
        """
        计算融合提示的置信度
        
        Args:
            fused_features: 融合后的特征 (B, prompt_dim)
            
        Returns:
            置信度分数 (B, 1)
        """
        return self.confidence_estimator(fused_features)
    
    def forward(self,
               memory_prompts: Optional[Dict[str, torch.Tensor]],
               periodic_prompts: Optional[Dict[str, torch.Tensor]],
               quality_prompts: Optional[Dict[str, torch.Tensor]],
               image_shape: Tuple[int, int]) -> Dict[str, torch.Tensor]:
        """
        融合三个模块的提示
        
        Args:
            memory_prompts: 记忆模块提示
            periodic_prompts: 周期模块提示
            quality_prompts: 质量模块提示
            image_shape: 图像形状 (H, W)
            
        Returns:
            融合后的掩码提示和置信度
        """
        # 提取特征
        memory_features = None
        if memory_prompts and 'memory_features' in memory_prompts:
            memory_features = memory_prompts['memory_features']
        elif memory_prompts and 'masks' in memory_prompts:
            # 从掩码中提取特征
            memory_features = memory_prompts['masks'].mean(dim=[-2, -1])
        
        periodic_features = None
        if periodic_prompts and 'phase_info' in periodic_prompts:
            periodic_features = periodic_prompts['phase_info'][:, -1, :]  # 取最后一个时间步
        elif periodic_prompts and 'masks' in periodic_prompts:
            periodic_features = periodic_prompts['masks'][:, -1, :, :].mean(dim=[-2, -1])
        
        quality_features = None
        if quality_prompts and 'quality_scores' in quality_prompts:
            quality_features = quality_prompts['quality_scores']
        elif quality_prompts and 'masks' in quality_prompts:
            quality_features = quality_prompts['masks'].mean(dim=[-2, -1])
        
        # 融合特征
        fused_features = self.fuse_features(
            memory_features, periodic_features, quality_features
        )
        
        # 生成掩码提示
        mask_prompt = self.generate_mask_prompt(fused_features, image_shape)
        
        # 计算置信度
        confidence = self.compute_confidence(fused_features)
        
        return {
            'mask_inputs': mask_prompt,
            'confidence': confidence,
            'fused_features': fused_features,
            'memory_features': memory_features,
            'periodic_features': periodic_features,
            'quality_features': quality_features
        }
